import re
import tkinter as tk
from tkinter import filedialog, messagebox, scrolledtext
import uuid
import os
from datetime import datetime
from bs4 import BeautifulSoup
from dashscope import api_key
from langchain.text_splitter import RecursiveCharacterTextSplitter, NLTKTextSplitter, SpacyTextSplitter
import markdownify
import PyPDF2
from docx import Document
import threading
from tkinter import ttk

from backend.qa_model import VectorSearchEngine  # 更新导入
from backend.updateToDatabase import VBAFile2Database

import spacy

class VBAHelperApp:
    def __init__(self, root):
        self.root = root
        self.root.title("经纬恒润VBA问答助手")
        self.root.geometry("400x300")

        # 主界面布局
        self.label = tk.Label(root, text="欢迎使用经纬恒润VBA问答助手", font=("Arial", 14))
        self.label.pack(pady=20)

        # 上传文件按钮
        self.upload_button = tk.Button(root, text="上传文件", command=self.start_upload_thread, width=20, height=2)
        self.upload_button.pack(pady=10)

        # 进度条
        self.progress = ttk.Progressbar(root, orient="horizontal", length=200, mode="indeterminate")
        self.progress.pack(pady=10)
        self.progress.pack_forget()  # Hide initially

        self.text_splitter = SpacyTextSplitter(
            chunk_size=256,  # 每块最大字符数
            chunk_overlap=40,  # 块间重叠字符数
            pipeline="zh_core_web_sm"  # 显式指定中文模型
        )
        #self.model_dir = snapshot_download('Jerry0/text2vec-base-chinese')
        #self.scentence_model = SentenceTransformer(self.model_dir)  # 英文轻量模型

        # 开始对话按钮
        self.chat_button = tk.Button(root, text="开始对话", command=self.start_chat, width=20, height=2)
        self.chat_button.pack(pady=10)

        self.DB_CONFIG = {
            "host": "hadoop102",
            "port": 15432,
            "database": "hirain_vba_usage",
            "user": "postgres",  # 替换实际数据库凭证
            "password": "postgres"
        }

    def start_upload_thread(self):
        """Start file upload in a background thread."""
        self.progress.pack()  # Show progress bar
        self.progress.start()  # Start animation
        self.upload_button.config(state=tk.DISABLED)  # Disable button during upload

        # Start the upload in a background thread
        thread = threading.Thread(target=self.upload_and_parse_file, daemon=True)
        thread.start()

    def upload_and_parse_file(self):
        """Process files in the background."""
        file_paths = filedialog.askopenfilenames(
            title="选择文件",
            filetypes=[("Markdown 文件", "*.md"), ("HTML 文件", "*.html"), ("PDF 文件", "*.pdf"), ("Word 文件", "*.docx"), ("所有文件", "*.*")]
        )
        if file_paths:
            folder_texts = []
            collection = file_paths[0].split("/")[-2]
            relation_text = []

            for file_path in file_paths:
                try:
                    if file_path.endswith(".pdf"):
                        text = self.parse_pdf(file_path)
                    elif file_path.endswith(".docx"):
                        text = self.parse_word(file_path)
                    elif file_path.endswith(".html"):
                        text = self.parse_html(file_path)
                    elif file_path.endswith(".md"):
                        text = self.parse_markdown(file_path)
                    else:
                        messagebox.showwarning("不支持的文件类型", "请上传 PDF、Word、HTML 或 Markdown 文件。")
                        continue
                    print(text)

                    #connection_text = file_path.split("/")[-1].split(" ")[-1].replace(".html", "").replace(".md", "") + ":: " + text
                    #folder_texts.append(connection_text)
                    relation = file_path.split("/")[-1].split(" ")[-1].replace(".html", "").replace(".md", "")

                    for item in text:
                        content = item["answer"]
                        question = item["question"]
                        docs = self.text_splitter.split_text(content)

                        docs = [{
                            'answer': doc,
                            'source': question,
                            'collection': collection
                        } for doc in docs]

                        uploader = VBAFile2Database(api_key="sk-72a6a7a8107d4bf2b80a90dae5bd1f8a")
                        uploader.store_in_db(docs, file_path)

                        if hasattr(self, "chat_text"):
                            self.add_message_to_chat(f"文件内容:\n{folder_texts}", is_user=False)
                except Exception as e:
                    messagebox.showerror("解析失败", f"解析文件时出错: {e}")

            self.root.after(0, lambda: messagebox.showinfo("解析成功", f"文件已解析: {os.path.basename(file_paths[0].split('/')[-2])}"))
            print(f"文件已解析: {os.path.basename(file_paths[0].split('/')[-2])}")
        self.root.after(0, self.reset_upload_ui)

    def reset_upload_ui(self):
        """Reset UI after upload completes."""
        self.progress.stop()
        self.progress.pack_forget()  # Hide progress bar
        self.upload_button.config(state=tk.NORMAL)  # Re-enable button

    def parse_pdf(self, file_path):
        text = ""
        with open(file_path, "rb") as file:
            reader = PyPDF2.PdfReader(file)
            for page in reader.pages:
                text += page.extract_text() + "\n"
        return text

    def parse_word(self, file_path):
        doc = Document(file_path)
        text = ""
        for para in doc.paragraphs:
            text += para.text + "\n"
        return text

    def parse_html(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8") as file:
                soup = BeautifulSoup(file.read(), "html.parser")
                
                # 查找特定的div元素
                main_div = soup.find("div", {
                    "class": "page-wrapper",
                    "role": "main",
                    "tabindex": "-1"
                })
                
                if not main_div:
                    return "未找到指定的页面内容区域"
                
                # 移除脚本和样式标签
                for script in main_div(["script", "style"]):
                    script.decompose()

                text = markdownify.markdownify(str(main_div))
                # 提取所有文本内容并合并为一段话
                # text = main_div.get_text(separator=" ", strip=True)
                
                # 去除不需要的文本内容
                # Remove unwanted texts
                unwanted_texts = [
                    "北京经纬恒润股份有限公司",
                    "VBA V3.0.0R",
                    "修订时间：",
                    "2025-01-22 16:07:29",
                    'results matching " "',
                    'No results matching " "',
                    'results matching ""',
                    'No results matching ""',
                    'No',
                    '===================',
                    '*',

                ]
                for unwanted in unwanted_texts:
                    text = text.replace(unwanted, "")

                    # Remove Markdown image syntax (e.g., ![alt](path))
                    text = re.sub(r'!\[.*?\]\(.*?\)', '', text)

                    # Normalize newlines and remove excessive whitespace
                    text = re.sub(r'\n+', '\n', text).strip()  # Replace multiple newlines with single newline

                    # Split by headings (e.g., #, ##) or bold text (**)
                    sections = []
                    current_section = {"title": "", "content": ""}
                    lines = text.split('\n')

                    for line in lines:
                        line = line.strip()
                        if line.startswith('#') or line.startswith('##') or line.startswith('###'):
                            # Save previous section if it has content
                            if current_section["content"]:
                                sections.append(current_section)
                                current_section = {"title": "", "content": ""}
                            # Extract heading as title
                            current_section["title"] = line.lstrip('#').strip()
                        #elif line.startswith('**') and line.endswith('**'):
                            # Use bold text as title if no heading exists
                            #if not current_section["title"]:
                                #current_section["title"] = line.strip('*').strip()
                        else:
                            # Add non-heading lines to content (skip empty lines)
                            if line:
                                current_section["content"] += line + ' '  # Replace newline with space

                    # Add the last section
                    if current_section["content"]:
                        sections.append(current_section)

                    # Format as structured data (e.g., list of Q&A pairs)
                    if sections:
                        qa_pairs = []
                        for section in sections:
                            if section["title"] and section["content"]:
                                # Clean up content (remove extra spaces from line joins)
                                cleaned_content = ' '.join(section["content"].split())
                                qa_pairs.append({
                                    "question": section["title"].split(" ")[-1],
                                    "answer": cleaned_content.replace("\n", "")
                                })
                        return qa_pairs if qa_pairs else [
                            {"question": file_path.split("/")[-1].split(" ")[-1].replace(".md", ""),
                             "answer": text.replace("\n", "")}]  # Fallback to raw text if no structure found
                    else:
                        return [{"question": file_path.split("/")[-1].split(" ")[-1].replace(".md", ""),
                                 "answer": text.replace("\n", "")}]  # Return raw text if no sections found

                # 返回单一段落文本
                # return text
                
        except Exception as e:
            return f"解析HTML文件时出错: {str(e)}"

    def parse_markdown(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8") as file:
                text = file.read()
                
                # Remove unwanted texts
                unwanted_texts = [
                    "北京经纬恒润股份有限公司",
                    "VBA V3.0.0R",
                    "修订时间：",
                    "2025-01-22 16:07:29",
                    'results matching " "',
                    'No results matching " "',
                    'results matching ""',
                    'No results matching ""',
                    'No',
                    '===================',
                    '*',

                ]
                for unwanted in unwanted_texts:
                    text = text.replace(unwanted, "")

                # Remove Markdown image syntax (e.g., ![alt](path))
                text = re.sub(r'!\[.*?\]\(.*?\)', '', text)

                # Normalize newlines and remove excessive whitespace
                text = re.sub(r'\n+', '\n', text).strip()  # Replace multiple newlines with single newline

                # Split by headings (e.g., #, ##) or bold text (**)
                sections = []
                current_section = {"title": "", "content": ""}
                lines = text.split('\n')
                
                for line in lines:
                    line = line.strip()
                    if line.startswith('#') or line.startswith('##') or line.startswith('###'):
                        # Save previous section if it has content
                        if current_section["content"]:
                            sections.append(current_section)
                            current_section = {"title": "", "content": ""}
                        # Extract heading as title
                        current_section["title"] = line.lstrip('#').strip()
                    #elif line.startswith('**') and line.endswith('**'):
                        # Use bold text as title if no heading exists
                        #if not current_section["title"]:
                            #current_section["title"] = line.strip('*').strip()
                    else:
                        # Add non-heading lines to content (skip empty lines)
                        if line:
                            current_section["content"] += line + ' '  # Replace newline with space

                # Add the last section
                if current_section["content"]:
                    sections.append(current_section)

                # Format as structured data (e.g., list of Q&A pairs)
                if sections:
                    qa_pairs = []
                    for section in sections:
                        if section["title"] and section["content"]:
                            # Clean up content (remove extra spaces from line joins)
                            cleaned_content = ' '.join(section["content"].split())
                            qa_pairs.append({
                                "question": section["title"].split(" ")[-1],
                                "answer": cleaned_content.replace("\n", "")
                            })
                    return qa_pairs if qa_pairs else [{"question": file_path.split("/")[-1].split(" ")[-1].replace(".md", ""), "answer": text.replace("\n", "")}]  # Fallback to raw text if no structure found
                else:
                    return [{"question": file_path.split("/")[-1].split(" ")[-1].replace(".md", ""), "answer": text.replace("\n", "")}]  # Return raw text if no sections found

        except Exception as e:
            return f"解析Markdown文件时出错: {str(e)}"

    def start_chat(self):
        # 关闭当前窗口
        self.root.destroy()
        # 创建对话窗口
        chat_window = tk.Tk()
        chat_window.title("对话界面")
        chat_window.geometry("800x500")

        # 存储会话数据
        self.sessions = {}  # {session_id: {"messages": [], "history": []}}
        self.current_session_id = None

        # 左边：聊天记录
        history_frame = tk.Frame(chat_window, width=200, height=500, bg="#f0f0f0")
        history_frame.pack(side=tk.LEFT, fill=tk.Y)

        # 添加新对话按钮
        new_chat_button = tk.Button(history_frame, text="新建对话", command=lambda: self.create_new_session(chat_window), width=20)
        new_chat_button.pack(pady=10)

        # 返回主界面按钮
        back_button = tk.Button(history_frame, text="返回主界面", command=lambda: self.return_to_main(chat_window), width=20)
        back_button.pack(pady=10)

        # 会话列表
        self.session_listbox = tk.Listbox(history_frame, width=25, height=30)
        self.session_listbox.pack(fill=tk.BOTH, expand=True)
        self.session_listbox.bind("<Double-Button-1>", lambda e: self.load_session(chat_window))

        # 右边：聊天框和输入框
        right_frame = tk.Frame(chat_window)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

        # 右上：聊天框（使用 Canvas 实现气泡效果）
        chat_frame = tk.Frame(right_frame, height=400, bg="#f5f5f5")
        chat_frame.pack(fill=tk.BOTH, expand=True)

        # 创建 Canvas 和 Scrollbar
        self.chat_canvas = tk.Canvas(chat_frame, bg="#f5f5f5", width=600, highlightthickness=0)
        scrollbar = tk.Scrollbar(chat_frame, orient="vertical", command=self.chat_canvas.yview)
        self.chat_canvas.configure(yscrollcommand=scrollbar.set)

        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.chat_canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 内部 Frame 用于放置消息
        self.chat_inner_frame = tk.Frame(self.chat_canvas, bg="#f5f5f5", width=600)
        self.chat_canvas.create_window((0, 0), window=self.chat_inner_frame, anchor="nw")

        # 绑定 Canvas 大小变化事件
        self.chat_inner_frame.bind("<Configure>", lambda e: self.chat_canvas.configure(scrollregion=self.chat_canvas.bbox("all")))

        # 右下：输入框和发送按钮
        input_frame = tk.Frame(right_frame, height=100)
        input_frame.pack(fill=tk.X)
        self.input_entry = tk.Entry(input_frame, font=("Arial", 12))
        self.input_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5, pady=5)
        
        # 绑定回车键事件
        self.input_entry.bind("<Return>", lambda event: self.send_message(chat_window))
        
        send_button = tk.Button(input_frame, text="发送", command=lambda: self.send_message(chat_window), width=10)
        send_button.pack(side=tk.RIGHT, padx=5, pady=5)

        # 初始化一个默认会话
        self.create_new_session(chat_window)

    def create_new_session(self, chat_window):
        session_id = str(uuid.uuid4())
        self.sessions[session_id] = {"messages": [], "history": []}
        self.current_session_id = session_id
        self.session_listbox.insert(tk.END, f"会话 {len(self.sessions)}")
        self.update_chat_display(chat_window)

    def load_session(self, chat_window):
        selected_index = self.session_listbox.curselection()
        if selected_index:
            session_id = list(self.sessions.keys())[selected_index[0]]
            self.current_session_id = session_id
            self.update_chat_display(chat_window)

    def update_chat_display(self, chat_window):
        # 清空聊天框
        for widget in self.chat_inner_frame.winfo_children():
            widget.destroy()

        if self.current_session_id:
            for msg in self.sessions[self.current_session_id]["messages"]:
                if msg.startswith("你:"):
                    self.add_message_to_chat(msg.replace("你:", "").strip(), is_user=True)
                else:
                    self.add_message_to_chat(msg, is_user=False)

    def add_message_to_chat(self, message, is_user=True):
        # 消息框架
        message_frame = tk.Frame(self.chat_inner_frame, bg="#f5f5f5")
        message_frame.pack(fill=tk.X, pady=5)

        # 消息内容框架
        content_frame = tk.Frame(message_frame, bg="#f5f5f5")
        content_frame.pack(side=tk.RIGHT if is_user else tk.LEFT, padx=10)

        # 气泡背景
        bubble_bg = "#95ec69" if is_user else "#ffffff"
        bubble = tk.Frame(content_frame, bg=bubble_bg, padx=10, pady=5, bd=1, relief="solid")
        bubble.pack()

        # 消息文本
        text_label = tk.Label(
            bubble,
            text=message,
            font=("Arial", 12),
            bg=bubble_bg,
            fg="black",
            wraplength=300,
            justify="left"
        )
        text_label.pack()

        # 时间戳
        timestamp = datetime.now().strftime("%H:%M")
        time_label = tk.Label(
            content_frame,
            text=timestamp,
            font=("Arial", 8),
            fg="gray",
            bg="#f5f5f5"
        )
        time_label.pack(anchor="e" if is_user else "w")

        # 更新 Canvas 滚动区域
        self.chat_canvas.configure(scrollregion=self.chat_canvas.bbox("all"))
        self.chat_canvas.yview_moveto(1.0)

    def send_message(self, chat_window):
        message = self.input_entry.get()
        if message and self.current_session_id:
            try:
                # 添加用户消息到界面
                self.add_message_to_chat(message, is_user=True)
                self.input_entry.delete(0, tk.END)

                # 构建系统提示
                system_prompt = """你是一个专业的VBA助手，请基于给定的上下文信息回答用户的问题。
                如果无法从上下文中找到相关信息，请明确告知。
                回答要简洁、准确，并尽可能引用相关的具体内容。"""

                # 获取历史对话
                history = self.sessions[self.current_session_id].get("history", [])
                self.search_engine = VectorSearchEngine(self.DB_CONFIG)
                
                # 执行检索
                retrieval_results = self.search_engine.similarity_search(
                    query_text=message,
                    top_n=5,
                    ef_search=200,
                    generate_answer=False
                )

                if not retrieval_results:
                    response = "抱歉，我找不到相关的信息来回答这个问题。"
                else:
                    # 整理检索到的上下文
                    contexts = []
                    sources = []
                    for result in retrieval_results:
                        if isinstance(result, dict):
                            contexts.append(result.get("content", ""))
                            if "source" in result:
                                sources.append(result["source"])

                    context_text = "\n".join(contexts)

                    # 生成回答
                    response_dict = self.search_engine.similarity_search(
                        query_text=message,
                        top_n=5,
                        ef_search=200,
                        generate_answer=True
                    )
                    print(response_dict)

                    if isinstance(response_dict, dict) and 'answer' in response_dict:
                        response = response_dict['answer']
                        if sources:
                            response += f"\n\n参考来源：{', '.join([source['collection'] + '--' + source['source'] for source in sources])}"
                    else:
                        response = "抱歉，生成回答时出现错误。"

                # 添加回复到界面
                self.add_message_to_chat(response, is_user=False)

                # 更新会话历史
                self.sessions[self.current_session_id]["messages"].extend([
                    f"用户: {message}",
                    f"助手: {response}"
                ])
                self.sessions[self.current_session_id]["history"].append({
                    "role": "user",
                    "content": message
                })
                self.sessions[self.current_session_id]["history"].append({
                    "role": "assistant",
                    "content": response
                })

            except Exception as e:
                error_message = f"处理消息时出错: {str(e)}"
                self.add_message_to_chat(error_message, is_user=False)

    def return_to_main(self, chat_window):
        # 关闭当前聊天窗口
        chat_window.destroy()
        # 重新创建主界面
        root = tk.Tk()
        app = VBAHelperApp(root)
        root.mainloop()

if __name__ == "__main__":

    root = tk.Tk()
    app = VBAHelperApp(root)
    root.mainloop()
